﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Text;
using OpenTK;

namespace LasRendering
{
    // K-d Tree
    // http://en.wikipedia.org/wiki/Kd-tree
    // Constructor and Add elements methods are multi-threaded.
    // Find nearest N neighbors method is single-thread.
    public class KdTree
    {
        private class ComparerByX : IComparer<Vertex>
        {
            public int Compare(Vertex element1, Vertex element2)
            {
                return element1.Position.X.CompareTo(element2.Position.X);
            }
        }

        private class ComparerByY : IComparer<Vertex>
        {
            public int Compare(Vertex element1, Vertex element2)
            {
                return element1.Position.Y.CompareTo(element2.Position.Y);
            }
        }

        private class ComparerByZ : IComparer<Vertex>
        {
            public int Compare(Vertex element1, Vertex element2)
            {
                return element1.Position.Z.CompareTo(element2.Position.Z);
            }
        }

        public int Dimensionality
        {
            private set;
            get;
        }

        public KdTreeNode Root
        {
            private set;
            get;
        }

        public KdTree(List<Vertex> elements, int dimensionality)
        {
            this.Dimensionality = dimensionality;
            this.Root = CreateNode(elements, 0, elements.Count - 1, 0);
        }

        private KdTreeNode CreateNode(List<Vertex> elements, int startIndex, int endIndex, int depth)
        {
            // Select axis based on depth so that axis cycles through all valid values
            int axis = depth % this.Dimensionality;

            var count = endIndex - startIndex + 1;
            if (count == 0)
            {
                return null;
            }
            if (count == 1)
            {
                var element = elements[startIndex];
                return new KdTreeNode(element, VectorHelper.GetValueByAxis(element.Position, axis), null, null);
            }

            // Sort point list and choose median as pivot element
            elements.Sort(startIndex, count, GetComparerByAxis(axis));
            var medianIndex = startIndex + count / 2;
            var medianElement = elements[medianIndex];

            // Create node and construct subtrees
            KdTreeNode leftChild = CreateNode(elements, startIndex, medianIndex - 1, depth + 1);
            KdTreeNode rightChild = CreateNode(elements, medianIndex + 1, endIndex, depth + 1);
            return new KdTreeNode(medianElement, VectorHelper.GetValueByAxis(medianElement.Position, axis), leftChild, rightChild);
        }

        private static IComparer<Vertex> GetComparerByAxis(int axis)
        {
            IComparer<Vertex> comparer;
            switch (axis)
            {
                case 0:
                    comparer = new ComparerByX();
                    break;
                case 1:
                    comparer = new ComparerByY();
                    break;
                case 2:
                    comparer = new ComparerByZ();
                    break;
                default:
                    comparer = null;
                    Debug.Fail(null);
                    break;
            }
            return comparer;
        }

        public void Add(List<Vertex> elements)
        {
            this.Add(elements, 0, elements.Count - 1, this.Root, 0);
        }

        private void Add(List<Vertex> elements, int startIndex, int endIndex, KdTreeNode node, int depth)
        {
            int axis = depth % this.Dimensionality;

            var count = endIndex - startIndex + 1;
            if (count == 0)
            {
                return;
            }
            if (count == 1)
            {
                var element = elements[startIndex];
                this.Add(element, node, depth);
                return;
            }

            // Sort point list and choose median as pivot element
            IComparer<Vertex> comparer = GetComparerByAxis(axis);
            elements.Sort(startIndex, count, comparer);
            int index = elements.BinarySearch(startIndex, count, node.Element, GetComparerByAxis(axis));
            if (index < 0)
            {
                index = ~index;
            }

            int leftStartIndex = -1, leftEndIndex = -1;
            int rightStartIndex = -1, rightEndIndex = -1;
            if (index <= startIndex)
            {
                rightStartIndex = startIndex;
                rightEndIndex = endIndex;
            }
            else if (index > endIndex)
            {
                leftStartIndex = startIndex;
                leftEndIndex = endIndex;
            }
            else
            {
                leftStartIndex = startIndex;
                leftEndIndex = index - 1;
                rightStartIndex = index;
                rightEndIndex = endIndex;
            }

            if (leftStartIndex >= 0)
            {
                if (node.LeftChild != null)
                {
                    this.Add(elements, leftStartIndex, leftEndIndex, node.LeftChild, depth + 1);
                }
                else
                {
                    node.LeftChild = this.CreateNode(elements, leftStartIndex, leftEndIndex, depth + 1);
                }
            }
            if (rightStartIndex >= 0)
            {
                if (node.RightChild != null)
                {
                    this.Add(elements, rightStartIndex, rightEndIndex, node.RightChild, depth + 1);
                }
                else
                {
                    node.RightChild = this.CreateNode(elements, rightStartIndex, rightEndIndex, depth + 1);
                }
            }
        }

        public void Add(Vertex element)
        {
            this.Add(element, this.Root, 0);
        }

        private void Add(Vertex element, KdTreeNode node, int depth)
        {
            // Check if node should be added to left or right sub-tree of current node.
            var axis = depth % this.Dimensionality;
            var location = VectorHelper.GetValueByAxis(element.Position, axis);
            if (location <= node.Location)
            {
                if (node.LeftChild == null)
                {
                    node.LeftChild = new KdTreeNode(element, location, null, null);
                }
                else
                {
                    this.Add(element, node.LeftChild, depth + 1);
                }
            }
            else
            {
                if (node.RightChild == null)
                {
                    node.RightChild = new KdTreeNode(element, location, null, null);
                }
                else
                {
                    this.Add(element, node.RightChild, depth + 1);
                }
            }
        }

        private float ComputeLength(Vector3 v)
        {
            switch (this.Dimensionality)
            {
                case 1:
                    return Math.Abs(v.X);
                case 2:
                    return (float)Math.Sqrt(v.X * v.X + v.Y * v.Y);
                case 3:
                    return v.LengthFast;
                default:
                    Debug.Fail(null);
                    return 0;
            }
        }

        public List<Vertex> FindInRange(Vector3 position, float range)
        {
            var elementList = new List<Vertex>();
            FindInRange(this.Root, position, range, 0, elementList);
            return elementList;
        }

        private void FindInRange(KdTreeNode node, Vector3 position, float range, int depth, List<Vertex> elementList)
        {
            if (node == null)
            {
                return;
            }

            var axis = depth % this.Dimensionality;
            var distance = ComputeLength(node.Element.Position - position);
            var location = VectorHelper.GetValueByAxis(position, axis);

            if (distance != 0 && distance < range)
            {
                elementList.Add(node.Element);
            }

            KdTreeNode nearChildNode, farChildNode;
            if (location < node.Location)
            {
                nearChildNode = node.LeftChild;
                farChildNode = node.RightChild;
            }
            else
            {
                nearChildNode = node.RightChild;
                farChildNode = node.LeftChild;
            }

            if (nearChildNode != null)
            {
                FindInRange(nearChildNode, position, range, depth + 1, elementList);
            }

            // The algorithm checks whether there could be any points on the other side of the splitting plane.
            if (Math.Abs(node.Location - location) < range)
            {
                if (farChildNode != null)
                {
                    FindInRange(farChildNode, position, range, depth + 1, elementList);
                }
            }
        }

        public PriorityQueue<float, Vertex> FindNearestNNeighbors(Vector3 position, int numNeighbors)
        {
            return FindNearestNNeighbors(position, numNeighbors, float.MaxValue);
        }

        public PriorityQueue<float, Vertex> FindNearestNNeighbors(Vector3 position, int numNeighbors, float range)
        {
            // Starting with the root node, the algorithm moves down the tree recursively, 
            // in the same way that it would if the search point were being inserted 
            // (i.e. it goes left or right depending on whether the point is less than or greater than the current node in the split dimension).
            var elementList = new PriorityQueue<float, Vertex>(numNeighbors);
            FindNearestNNeighbors(this.Root, position, ref range, numNeighbors, 0, elementList);
            return elementList;
        }

        private void FindNearestNNeighbors(KdTreeNode node, Vector3 position, ref float currentNBest, int numNeighbors, int depth, PriorityQueue<float, Vertex> elementList)
        {
            if (node == null)
            {
                return;
            }

            var axis = depth % this.Dimensionality;
            var distance = this.ComputeLength(node.Element.Position - position);
            var location = VectorHelper.GetValueByAxis(position, axis);

            // If the current node is closer than the current best, then it becomes the current best.
            if (distance != 0 && distance < currentNBest)
            {
                if (elementList.Count == numNeighbors)
                {
                    elementList.Pop();
                }
                elementList.Push(distance, node.Element);

                if (elementList.Count == numNeighbors)
                {
                    currentNBest = elementList.MaxPriority;
                }
            }

            // The algorithm checks whether there could be any points on the other side of the splitting plane 
            // that are closer to the search point than the current best. 
            // In concept, this is done by intersecting the splitting hyperplane with a hypersphere around the search point 
            // that has a radius equal to the current nearest distance. Since the hyperplanes are all axis-aligned 
            // this is implemented as a simple comparison to see whether the difference between the splitting coordinate of 
            // the search point and current node is less than the distance (overall coordinates) from the search point to the current best. 

            KdTreeNode nearChildNode, farChildNode;
            if (location < node.Location)
            {
                nearChildNode = node.LeftChild;
                farChildNode = node.RightChild;
            }
            else
            {
                nearChildNode = node.RightChild;
                farChildNode = node.LeftChild;
            }

            // If the hypersphere doesn't intersect the splitting plane, then the algorithm continues walking up the tree, 
            // and the entire branch on the other side of that node is eliminated.
            if (nearChildNode != null)
            {
                FindNearestNNeighbors(nearChildNode, position, ref currentNBest, numNeighbors, depth + 1, elementList);
            }

            // If the hypersphere crosses the plane, there could be nearer points on the other side of the plane, 
            // so the algorithm must move down the other branch of the tree from the current node looking for closer points, 
            // following the same recursive process as the entire search.
            if (Math.Abs(node.Location - location) < currentNBest)
            {
                if (farChildNode != null)
                {
                    FindNearestNNeighbors(farChildNode, position, ref currentNBest, numNeighbors, depth + 1, elementList);
                }
            }
        }
    }
}
